#!/usr/bin/python3
#   Main program of pg_copyplot.
#
#   Copyright (C) 2017 Markus Zeindl
#
#   This file is part of pg_copyplot, an adapter for plotting tools 
#   employing postgresql's COPY command.
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.

import argparse
import collections
import csv
import io
import itertools
import os
import subprocess
import sys
import tempfile

class CommandlineArgument:
    def __init__(self, switch, caption, type, default, helptext):
        self._switch = switch
        self._name = caption
        self._type = type
        self._default = default
        self._helptext = helptext

    @property
    def name(self):
        return self._name

    @property
    def type(self):
        return self._type

    @property
    def default(self):
        return self._default

    @property
    def helptext(self):
        return self._helptext

    def add_to_argparse_parser(self, parser_object):
        parser_object.add_argument("-" + self._switch,
                                   metavar=self._name,
                                   type=self._type,
                                   default=self._default,
                                   help=self._helptext) 

    def evaluate_from_argparse_args(self, args_obj):
        return args_obj.__dict__[ self._switch ]

    
class Commandline(collections.UserDict):
    ARG_SOURCE = CommandlineArgument("s", "INPUT", argparse.FileType('r'), sys.stdin, "specify file to read from")
    ARG_XSPEC = CommandlineArgument("x", "COLUMN_INDEX_OF_X", int, 1, "specify index of x values")
    ARG_YSPEC = CommandlineArgument("y", "COLUMN_INDEX_OF_Y", int, 1, "specify index of y values")
    ARG_COLUMNDELIM = CommandlineArgument("d", "COLUMN_DELIMITER", str, ",", "specify delimiter of columns")
    ARG_DESTINATION_PORT = CommandlineArgument("p", "DESTINATION_PORT", int, 9000, "specify nc listen port")
    ARG_FACTORSPEC = CommandlineArgument("f", "COLUMN_INDEX_OF_FACTOR", int, 0, "specify one-based index of factor levels, (defaults to None)")
    
    params = [ ARG_SOURCE,
               ARG_XSPEC,
               ARG_YSPEC,
               ARG_COLUMNDELIM,
               ARG_DESTINATION_PORT,
               ARG_FACTORSPEC ]
    
    @staticmethod
    def get_defaults():
        defaults = dict()
        for arg in Commandline.params:
            defaults[arg.name] = arg.default
        return defaults
    
    def parse(self, arguments=sys.argv[1:]):
        parser = argparse.ArgumentParser()
        self.add_arguments(parser)
        args = parser.parse_args(args=arguments)
        self.evaluate_args(args)

    def add_arguments(self, parser_obj):
        for param in Commandline.params:
            param.add_to_argparse_parser(parser_obj)

    def evaluate_args(self, args):
        for param in Commandline.params:
            self[param.name] = param.evaluate_from_argparse_args(args)

ScatterPlotSeries = collections.namedtuple("ScatterPlotSeries", [ "datafile", "delimiter", "index_x", "index_y" ])

class ScatterPlot:
    def __init__(self, gnuplot_instance):
        self.gnu_plot = gnuplot_instance
        
    def plot(self, datafile, delimiter, index_x, index_y):
        gnuplot_command = "plot '{}' using {}:{} with linespoints".format(datafile,
                                                                      index_x,
                                                                      index_y)
        self.gnu_plot.set_variable(SimpleGnuplot.DATAFILE_COLUMN_DELIMITER, delimiter)
        self.gnu_plot.add_statement(gnuplot_command)
        self.gnu_plot.execute()

    def plot_multiple_series(self, series_to_plot):
        series_tpl = "'{}' using {}:{} with linespoints"
        formatted_series = [series_tpl.format(single_series.datafile,
                                              single_series.index_x,
                                              single_series.index_y) for single_series in series_to_plot]
        command_tpl = "plot {}"
        command = command_tpl.format(", ".join(formatted_series))
        self.gnu_plot.set_variable(SimpleGnuplot.DATAFILE_COLUMN_DELIMITER, series_to_plot[0].delimiter)
        self.gnu_plot.add_statement(command)
        self.gnu_plot.execute()
            
class SimpleGnuplot:
    DATAFILE_COLUMN_DELIMITER = "datafile separator"
    DEFAULT_DELIMITER = '\t'
    
    def __init__(self, executable="/usr/bin/gnuplot", args=[ "-p" ]):
        self._executable = executable
        self._args = args
        self._gnuplot_program = []
        
    def set_variable(self, varname, value):
        if type(value) == str:
            value = "\"{}\"".format(value)
        
        self.add_statement("set {} {}".format(varname, value))

    def add_statement(self, statement):
        self._gnuplot_program.append(statement)

    def add_inlinedata_and_get_datafilehandle(self, id, data):
        datafilehandle = "${}".format(id)
        command_tpl = "{} <<EOF\n{}\nEOF"
        command = command_tpl.format(datafilehandle, data)
        self.add_statement(command)
        return datafilehandle
        
    def get_complete_program(self):
        return "\n".join(self._gnuplot_program)
        
    def execute(self):
        self.program_per_stdin_execute()

    def program_per_stdin_execute(self):
        command = [self._executable] + self._args
        process = subprocess.Popen(command,
                                   stdin=subprocess.PIPE,
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE)
        process.stdin.write(self.get_complete_program().encode())
        process.stdin.write("\n".encode())
        print(process.communicate()[1].decode())

    def tempfile_execute(self, method):
        tempfilename = self.create_and_prepare_tempfile()
        method(tempfilename)
        self.wipe_tempfile(tempfilename)
        
    def create_and_prepare_tempfile(self):
        (fd, filename) = tempfile.mkstemp()
        os.close(fd)
        buffer = open(filename, "w")
        buffer.write(self.get_complete_program())
        buffer.close()                
        return filename

    def wipe_tempfile(self, filename):
        os.unlink(filename)        

class DataSplitter:
    def __init__(self, data, column_delimiter):
        self.databuffer = io.StringIO(data)
        self.column_delimiter = column_delimiter

    def split_by_factor(self, factor_index):
        ret = collections.defaultdict(str)
        self.rows = self.parse_fields()
        keyfunc = lambda row: row[factor_index]

        self.rows.sort(key=keyfunc)
        for (field, subiter) in itertools.groupby(self.rows, keyfunc):
            datapartition = io.StringIO()
            writer = csv.writer(datapartition, delimiter=self.column_delimiter, lineterminator="\n")
            for m in subiter:
                writer.writerow(m)
            ret[field] += datapartition.getvalue()
        return dict(ret)

    def parse_fields(self):
        ret = []
        csv_reader = csv.reader(self.databuffer, delimiter=self.column_delimiter)
        for row in csv_reader:
            ret.append(row)
        return ret
            
def read_plot_data(stream):
    ret = stream.read()
    stream.close()
    return ret

def create_tempdatafile_with_plot_data_and_return_filename(plot_data):  
    (fd, filename) = tempfile.mkstemp()
    os.close(fd)
    tmpfile = open(filename, 'w')
    tmpfile.write(plot_data)
    tmpfile.write("\n")
    tmpfile.close()
    return filename

def produce_simple_scatterplot(c):
    plot_data = read_plot_data(c[Commandline.ARG_SOURCE.name])
    dataid = "M"
    
    gnuplot = SimpleGnuplot(executable="/bin/nc", args=[ "127.0.0.1", str(c[Commandline.ARG_DESTINATION_PORT.name]) ])
    datafilehandle = gnuplot.add_inlinedata_and_get_datafilehandle(dataid, plot_data)
    scatterplot = ScatterPlot(gnuplot)
    scatterplot.plot(datafilehandle,
                     c[Commandline.ARG_COLUMNDELIM.name],
                     c[Commandline.ARG_XSPEC.name],
                     c[Commandline.ARG_YSPEC.name])

def produce_factored_scatterplot(c):
    plot_data = read_plot_data(c[Commandline.ARG_SOURCE.name])
    ONEBASED_TO_ZEROBASED = -1
    ds = DataSplitter(plot_data, c[Commandline.ARG_COLUMNDELIM.name])
    splitted_data = ds.split_by_factor(c[Commandline.ARG_FACTORSPEC.name]+ONEBASED_TO_ZEROBASED)
    gnuplot = SimpleGnuplot(executable="/bin/nc", args=[ "127.0.0.1", str(c[Commandline.ARG_DESTINATION_PORT.name]) ])
    series = []
    for level in splitted_data:
        datafilehandle = gnuplot.add_inlinedata_and_get_datafilehandle(level.replace(" ", "_"), splitted_data[level])
        series.append(ScatterPlotSeries(datafile=datafilehandle,
                                        delimiter=c[Commandline.ARG_COLUMNDELIM.name],
                                        index_x=c[Commandline.ARG_XSPEC.name],
                                        index_y=c[Commandline.ARG_YSPEC.name]))
        
    scatterplot = ScatterPlot(gnuplot)
    scatterplot.plot_multiple_series(series)
    
    
if __name__ == "__main__":
    c = Commandline()
    c.parse()

    if c[Commandline.ARG_FACTORSPEC.name] == -1:
        produce_simple_scatterplot(c)
    elif c[Commandline.ARG_FACTORSPEC.name] > 0:
        produce_factored_scatterplot(c)
